{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial 06, case 5: Stokes problem with distributed control\n",
    "\n",
    "In this tutorial we solve the optimal control problem\n",
    "\n",
    "$$\\min J(y, u) = \\frac{1}{2} \\int_{\\Omega} |v - v_d|^2 dx + \\frac{\\alpha}{2} \\int_{\\Omega} |u|^2 dx$$\n",
    "s.t.\n",
    "$$\\begin{cases}\n",
    "- \\Delta v + \\nabla p = f + u   & \\text{in } \\Omega\\\\\n",
    "         \\text{div} v = 0       & \\text{in } \\Omega\\\\\n",
    "                    v = 0       & \\text{on } \\partial\\Omega\n",
    "\\end{cases}$$\n",
    "\n",
    "where\n",
    "$$\\begin{align*}\n",
    "& \\Omega                      & \\text{unit square}\\\\\n",
    "& u \\in [L^2(\\Omega)]^2       & \\text{control variable}\\\\\n",
    "& v \\in [H^1_0(\\Omega)]^2     & \\text{state velocity variable}\\\\\n",
    "& p \\in L^2(\\Omega)           & \\text{state pressure variable}\\\\\n",
    "& \\alpha > 0                  & \\text{penalization parameter}\\\\\n",
    "& v_d                         & \\text{desired state}\\\\\n",
    "& f                           & \\text{forcing term}\n",
    "\\end{align*}$$\n",
    "using an adjoint formulation solved by a one shot approach"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from mpi4py import MPI\n",
    "from petsc4py import PETSc\n",
    "import sympy\n",
    "from ufl import div, grad, inner, Measure, replace, TestFunction, TrialFunction\n",
    "from dolfinx import Constant, DirichletBC, Function, FunctionSpace, VectorFunctionSpace\n",
    "from dolfinx.fem import locate_dofs_topological\n",
    "from dolfinx.io import XDMFFile\n",
    "from dolfinx.plot import create_vtk_topology\n",
    "from multiphenicsx.fem import (assemble_matrix_block, assemble_scalar, assemble_vector_block,\n",
    "                               BlockVecSubVectorWrapper, create_vector_block)\n",
    "import pyvista"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Mesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with XDMFFile(MPI.COMM_WORLD, \"data/square.xdmf\", \"r\") as infile:\n",
    "    mesh = infile.read_mesh()\n",
    "    mesh.topology.create_connectivity_all()\n",
    "    subdomains = infile.read_meshtags(mesh, name=\"subdomains\")\n",
    "    boundaries = infile.read_meshtags(mesh, name=\"boundaries\")\n",
    "boundaries_1234 = boundaries.indices[np.isin(boundaries.values, (1, 2, 3, 4))]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define associated measures\n",
    "dx = Measure(\"dx\")(subdomain_data=subdomains)\n",
    "ds = Measure(\"ds\")(subdomain_data=boundaries)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dolfinx_to_pyvista_mesh(mesh):\n",
    "    num_cells = mesh.topology.index_map(mesh.topology.dim).size_local\n",
    "    cell_entities = np.arange(num_cells, dtype=np.int32)\n",
    "    pyvista_cells, cell_types = create_vtk_topology(mesh, mesh.topology.dim, cell_entities)\n",
    "    grid = pyvista.UnstructuredGrid(pyvista_cells, cell_types, mesh.geometry.x)\n",
    "    return grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pyvista_mesh_plot(mesh):\n",
    "    grid = dolfinx_to_pyvista_mesh(mesh)\n",
    "    plotter = pyvista.PlotterITK()\n",
    "    plotter.add_mesh(grid)\n",
    "    plotter.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_mesh_plot(mesh)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function spaces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Y_velocity = VectorFunctionSpace(mesh, (\"Lagrange\", 2))\n",
    "Y_pressure = FunctionSpace(mesh, (\"Lagrange\", 1))\n",
    "U = VectorFunctionSpace(mesh, (\"Lagrange\", 2))\n",
    "Q_velocity = Y_velocity.clone()\n",
    "Q_pressure = Y_pressure.clone()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Trial and test functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(v, p) = (TrialFunction(Y_velocity), TrialFunction(Y_pressure))\n",
    "(w, q) = (TestFunction(Y_velocity), TestFunction(Y_pressure))\n",
    "u = TrialFunction(U)\n",
    "r = TestFunction(U)\n",
    "(z, b) = (TrialFunction(Q_velocity), TrialFunction(Q_pressure))\n",
    "(s, d) = (TestFunction(Q_velocity), TestFunction(Q_pressure))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ### Problem data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha = 1.e-5\n",
    "epsilon = 1.e-5\n",
    "x, y = sympy.symbols(\"x[0], x[1]\")\n",
    "psi_d = 10 * (1 - sympy.cos(0.8 * np.pi * x)) * (1 - sympy.cos(0.8 * np.pi * y)) * (1 - x)**2 * (1 - y)**2\n",
    "v_d_x = sympy.lambdify([x, y], psi_d.diff(y, 1))\n",
    "v_d_y = sympy.lambdify([x, y], - psi_d.diff(x, 1))\n",
    "v_d = Function(Y_velocity)\n",
    "v_d.interpolate(lambda x: np.stack((v_d_x(x[0], x[1]), v_d_y(x[0], x[1])), axis=0))\n",
    "ff = Constant(mesh, (0., 0.))\n",
    "bc0 = Function(Y_velocity)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optimality conditions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = [[inner(v, w) * dx, None, None, inner(grad(z), grad(w)) * dx, - b * div(w) * dx],\n",
    "     [None, None, None, - q * div(z) * dx, epsilon * b * q * dx],\n",
    "     [None, None, alpha * inner(u, r) * dx, - inner(z, r) * dx, None],\n",
    "     [inner(grad(v), grad(s)) * dx, - p * div(s) * dx, - inner(u, s) * dx, None, None],\n",
    "     [- d * div(v) * dx, epsilon * p * d * dx, None, None, None]]\n",
    "f = [inner(v_d, w) * dx,\n",
    "     None,\n",
    "     None,\n",
    "     inner(ff, s) * dx,\n",
    "     None]\n",
    "a[3][3] = Constant(mesh, 0.) * inner(z, s) * dx\n",
    "f[1] = Constant(mesh, 0.) * q * dx\n",
    "f[2] = inner(Constant(mesh, (0., 0.)), r) * dx\n",
    "f[4] = Constant(mesh, 0.) * d * dx\n",
    "bdofs_Y_velocity_1234 = locate_dofs_topological((Y_velocity, Y_velocity), mesh.topology.dim - 1, boundaries_1234)\n",
    "bdofs_Q_velocity_1234 = locate_dofs_topological((Q_velocity, Y_velocity), mesh.topology.dim - 1, boundaries_1234)\n",
    "bc = [DirichletBC(bc0, bdofs_Y_velocity_1234, Y_velocity),\n",
    "      DirichletBC(bc0, bdofs_Q_velocity_1234, Q_velocity)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(v, p) = (Function(Y_velocity), Function(Y_pressure))\n",
    "u = Function(U)\n",
    "(z, b) = (Function(Q_velocity), Function(Q_pressure))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cost functional"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "J = 0.5 * inner(v - v_d, v - v_d) * dx + 0.5 * alpha * inner(u, u) * dx"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Uncontrolled functional value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract state forms from the optimality conditions\n",
    "a_state = [[replace(a[i][j], {s: w, d: q}) if a[i][j] is not None else None\n",
    "            for j in (0, 1)] for i in (3, 4)]\n",
    "f_state = [replace(f[i], {s: w, d: q}) for i in (3, 4)]\n",
    "bc_state = [bc[0]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assemble the block linear system for the state\n",
    "A_state = assemble_matrix_block(a_state, bcs=bc_state)\n",
    "A_state.assemble()\n",
    "F_state = assemble_vector_block(f_state, a_state, bcs=bc_state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solve\n",
    "vp = create_vector_block([f[j] for j in (0, 1)])\n",
    "ksp = PETSc.KSP()\n",
    "ksp.create(mesh.mpi_comm())\n",
    "ksp.setOperators(A_state)\n",
    "ksp.setType(\"preonly\")\n",
    "ksp.getPC().setType(\"lu\")\n",
    "ksp.getPC().setFactorSolverType(\"mumps\")\n",
    "ksp.setFromOptions()\n",
    "ksp.solve(F_state, vp)\n",
    "vp.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Split the block solution in components\n",
    "with BlockVecSubVectorWrapper(vp, [c.function_space.dofmap for c in (v, p)]) as vp_wrapper:\n",
    "    for vp_wrapper_local, component in zip(vp_wrapper, (v, p)):\n",
    "        with component.vector.localForm() as component_local:\n",
    "            component_local[:] = vp_wrapper_local"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "J_uncontrolled = mesh.mpi_comm().allreduce(assemble_scalar(J), op=MPI.SUM)\n",
    "print(\"Uncontrolled J =\", J_uncontrolled)\n",
    "assert np.isclose(J_uncontrolled, 0.1784536)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pyvista_scalar_field_plot(mesh, scalar_field, name):\n",
    "    grid = dolfinx_to_pyvista_mesh(mesh)\n",
    "    grid.point_arrays[name] = scalar_field.compute_point_values()\n",
    "    grid.set_active_scalars(name)\n",
    "    plotter = pyvista.PlotterITK()\n",
    "    plotter.add_mesh(grid)\n",
    "    plotter.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pyvista_vector_field_plot(mesh, vector_field, name, factor):\n",
    "    grid = dolfinx_to_pyvista_mesh(mesh)\n",
    "    values = np.zeros((mesh.geometry.x.shape[0], 3))\n",
    "    values[:, :2] = vector_field.compute_point_values()\n",
    "    grid.point_arrays[name] = values\n",
    "    grid.set_active_vectors(name)\n",
    "    plotter = pyvista.PlotterITK()\n",
    "    plotter.add_mesh(grid)\n",
    "    glyphs = grid.glyph(orient=name, factor=factor)\n",
    "    plotter.add_mesh(glyphs)\n",
    "    plotter.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_vector_field_plot(mesh, v, \"uncontrolled state velocity\", factor=1e-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_scalar_field_plot(mesh, p, \"uncontrolled state pressure\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optimal control"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assemble the block linear system for the optimality conditions\n",
    "A = assemble_matrix_block(a, bcs=bc)\n",
    "A.assemble()\n",
    "F = assemble_vector_block(f, a, bcs=bc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solve\n",
    "vpuzb = create_vector_block(f)\n",
    "ksp = PETSc.KSP()\n",
    "ksp.create(mesh.mpi_comm())\n",
    "ksp.setOperators(A)\n",
    "ksp.setType(\"preonly\")\n",
    "ksp.getPC().setType(\"lu\")\n",
    "ksp.getPC().setFactorSolverType(\"mumps\")\n",
    "ksp.setFromOptions()\n",
    "ksp.solve(F, vpuzb)\n",
    "vpuzb.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Split the block solution in components\n",
    "with BlockVecSubVectorWrapper(vpuzb, [c.function_space.dofmap for c in (v, p, u, z, b)]) as vpuzb_wrapper:\n",
    "    for vpuzb_wrapper_local, component in zip(vpuzb_wrapper, (v, p, u, z, b)):\n",
    "        with component.vector.localForm() as component_local:\n",
    "            component_local[:] = vpuzb_wrapper_local"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "J_controlled = mesh.mpi_comm().allreduce(assemble_scalar(J), op=MPI.SUM)\n",
    "print(\"Optimal J =\", J_controlled)\n",
    "assert np.isclose(J_controlled, 0.0052941)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_vector_field_plot(mesh, v, \"state velocity\", factor=1e-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_scalar_field_plot(mesh, p, \"state pressure\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_vector_field_plot(mesh, u, \"control\", factor=1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_vector_field_plot(mesh, z, \"adjoint velocity\", factor=1e2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_scalar_field_plot(mesh, b, \"adjoint pressure\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython"
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
